import os
import numpy as np
import logging
import hashlib
from PIL import Image

# 配置日志
logger = logging.getLogger(__name__)

class CacheModule:
    def __init__(self, cache_dir, dataset_path):
        self.cache_dir = cache_dir
        self.dataset_path = dataset_path
        self.dataset_features = {}
        self.path_hash_mapping = {}

    def get_feature_cache_path(self, image_path):
        """获取特征缓存文件路径"""
        # 使用相对路径的哈希值作为缓存文件名，避免路径过长
        rel_path = os.path.relpath(image_path, self.dataset_path)
        image_hash = hashlib.md5(rel_path.encode()).hexdigest()
        cache_path = os.path.join(self.cache_dir, f"{image_hash}.npy")
        return cache_path, image_hash

    def save_feature_cache(self, image_path, features):
        """保存特征到缓存文件"""
        try:
            if not os.path.exists(self.cache_dir):
                os.makedirs(self.cache_dir)
            cache_path, _ = self.get_feature_cache_path(image_path)
            np.save(cache_path, features)
            logger.info(f"特征缓存已保存: {cache_path}")
        except Exception as e:
            logger.error(f'保存缓存失败 {image_path}: {str(e)}')

    def load_feature_cache(self, image_path):
        """从缓存文件加载特征"""
        try:
            cache_path, _ = self.get_feature_cache_path(image_path)
            if os.path.exists(cache_path):
                return np.load(cache_path)
        except Exception as e:
            logger.error(f'加载缓存失败 {image_path}: {str(e)}')
        return None

    def load_all_features(self):
        """从缓存加载所有数据集特征"""
        if os.path.exists(self.cache_dir):
            try:
                # 首先遍历数据集获取所有图片文件路径
                image_paths = []
                for root, _, files in os.walk(self.dataset_path):
                    for file in files:
                        if file.lower().endswith(('.png', '.jpg', '.jpeg')):
                            image_paths.append(os.path.join(root, file))
                
                # 加载每个图片对应的特征
                for image_path in image_paths:
                    try:
                        cache_path, file_hash = self.get_feature_cache_path(image_path)
                        if os.path.exists(cache_path):
                            features = np.load(cache_path)
                            self.dataset_features[image_path] = features
                            self.path_hash_mapping[file_hash] = image_path
                    except Exception as e:
                        logger.error(f'加载特征缓存出错 {image_path}: {str(e)}')
                        continue
                
                logger.info(f"已加载特征文件：{len(self.dataset_features)} 个")
            except Exception as e:
                logger.error(f'加载特征缓存目录出错: {str(e)}')
        return self.dataset_features

    def get_features(self):
        """获取所有特征"""
        return self.dataset_features

    def clear_cache(self):
        """清空缓存"""
        try:
            if os.path.exists(self.cache_dir):
                for file in os.listdir(self.cache_dir):
                    file_path = os.path.join(self.cache_dir, file)
                    try:
                        if os.path.isfile(file_path):
                            os.unlink(file_path)
                    except Exception as e:
                        logger.error(f'删除缓存文件失败 {file_path}: {str(e)}')
            self.dataset_features.clear()
            self.path_hash_mapping.clear()
            logger.info("缓存已清空")
        except Exception as e:
            logger.error(f'清空缓存失败: {str(e)}')